import os
import operator
from dotenv import load_dotenv
from typing import Annotated, TypedDict, Union
from langchain_core.messages import BaseMessage, HumanMessage, ToolMessage
from langchain_community.tools import DuckDuckGoSearchRun
from langchain_core.utils.function_calling import convert_to_openai_tool
from langchain_openai import ChatOpenAI

from langgraph.graph import StateGraph, END
from langgraph.prebuilt import ToolNode

# --- 1. 加载环境变量 ---
load_dotenv()

# --- 2. 定义工具 ---
# 我们创建一个 DuckDuckGo 搜索工具实例。
# LangChain 工具默认包含详细的描述，LLM 会利用这些描述来决定何时以及如何使用工具。
search_tool = DuckDuckGoSearchRun()
tools = [search_tool]

# --- 3. 定义状态 ---
# 这是一个特殊的状态，用于处理对话历史。
# `operator.add` 意味着每当 "messages" 被更新时，新的消息会追加到现有列表的末尾，
# 而不是替换它。这对于维护对话历史至关重要。
class AgentState(TypedDict):
    messages: Annotated[list, operator.add]

# --- 4. 定义图中的节点 ---

# 初始化模型并绑定工具。
# `.bind_tools()` 会将工具的定义（包括名称、参数、描述）转换为 OpenAI 函数调用格式，
# 并告知模型这些工具是可用的。
llm = ChatOpenAI(model="qwen-plus-latest", base_url=os.getenv("OPENAI_BASE_URL"))
model_with_tools = llm.bind_tools(tools)

# a. Agent 节点
# 这个节点负责调用 LLM。它是代理的核心 "大脑"。
def agent_node(state: AgentState):
    """
    调用 LLM 来决定下一步行动。

    Args:
        state (AgentState): 图的当前状态。

    Returns:
        dict: 包含 LLM 响应的消息。
    """
    print("---AGENT: 思考中...---")
    response = model_with_tools.invoke(state["messages"])
    # 我们将模型的响应（可能是一个普通回答，也可能是一个工具调用请求）返回，
    # 它将被添加到状态的 messages 列表中。
    return {"messages": [response]}

# b. Tool 节点
# 这个节点负责执行工具。它是一个"执行者"。
# `ToolNode` 是 langgraph 提供的一个预构建节点，能简化工具的调用过程。
# 它会自动解析模型返回的 ToolMessage，执行对应的工具，并将结果格式化为 ToolMessage 返回。
tool_node = ToolNode(tools)

# c. 最终节点
def final_answer_node(state: AgentState):
    """
    一个简单的节点，标志着流程的结束。
    """
    print("---流程结束---")
    return {}

# --- 5. 定义条件边 ---
# 这个函数决定在 agent 节点之后应该做什么。
def should_continue(state: AgentState) -> str:
    """
    根据 Agent 的最新响应决定工作流的下一站。

    Args:
        state (AgentState): 图的当前状态。

    Returns:
        str: 'tools' 或 '__end__'。
    """
    print("---ROUTER: 决策中...---")
    last_message = state["messages"][-1]

    # 如果 `tool_calls` 字段中没有任何内容，意味着 Agent 认为它已经完成了任务，
    # 可以直接向用户返回最终答案了。
    if not last_message.tool_calls:
        print("---ROUTER: 任务完成，结束流程。---")
        return "__end__"
    # 否则，意味着 Agent 希望调用一个或多个工具来获取更多信息。
    else:
        print("---ROUTER: 需要调用工具。---")
        return "tools" # 这是我们给 ToolNode 起的名字

# --- 6. 构建图 ---
workflow = StateGraph(AgentState)

# a. 添加节点
workflow.add_node("agent", agent_node)
workflow.add_node("tools", tool_node)

# b. 设置入口点
workflow.set_entry_point("agent")

# c. 添加条件边
workflow.add_conditional_edges(
    "agent",          # 决策点：在 agent 节点执行完毕后
    should_continue,  # 决策函数：使用此函数来决定去向
)

# d. 添加普通边
# 在工具节点执行完毕后，我们总是希望将结果返回给 agent 节点，
# 让它根据工具的输出继续思考。这就形成了一个循环。
workflow.add_edge("tools", "agent")

# e. 编译图
app = workflow.compile()

# --- 7. 运行图 ---
if __name__ == "__main__":
    # 我们用一个需要实时信息的问题来测试它
    question = "通义千问的最新版本是什么？"
    
    # 将用户问题包装成 HumanMessage，这是 langgraph 对话状态的标准格式
    inputs = {"messages": [HumanMessage(content=question)]}
    
    # 流式输出，这样我们可以清晰地看到每一步的执行过程
    for event in app.stream(inputs, stream_mode="values"):
        # `event` 是完整的状态快照
        event["messages"][-1].pretty_print()
        print("\n" + "="*50 + "\n") 